package mo.offer_version2;

import java.util.Random;

public class L40 {

    private Random random = new Random();

    public int[] getLeastNumbers(int[] arr, int k) {
        randomSelect(arr, 0, arr.length-1, k);
        int[] ret = new int[k];
        for (int i = 0; i < k; i++) {
            ret[i] = arr[i];
        }
        return ret;
    }

    private void randomSelect(int[] arr, int l, int r, int k) {
        if (l >= r) {
            return;
        }
        int pos = randomPartition(arr, l, r);
        int num = pos - l + 1;
        if (k == num) {
            return;
        } else if (k < num) {
            randomSelect(arr, l, pos-1, k);
        } else {
            randomSelect(arr, pos+1, r, k-num);
        }
    }

    private int randomPartition(int[] arr, int l, int r) {
        int i = random.nextInt(r-l+1) + l;
        swap(arr, i, r);
        return partition(arr, l, r);
    }

    private int partition(int[] nums, int l, int r) {
        int pivot = nums[r];
        int i = l - 1;
        for (int j = l; j <= r-1; j++) {
            if (nums[j] <= pivot) {
                i = i + 1;
                swap(nums, i, j);
            }
        }
        swap(nums, i+1, r);
        return i+1;
    }

    private void swap(int[] nums, int i, int j) {
        if (i == j) {
            return;
        }
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }



}
